from gurobipy import *
import numpy as np
from traffic_engineer import *
from traffic_engineering_class import *

class DualHullTrafficEngineer(TrafficEngineer):
	## Takes in two new variables:
	## update_frequency - how much time in-between every traffic engineering events
	def __init__(self, aurora_network, all_interblock_paths, update_frequency, training_length, reduce_multihop=False, ignore_flows=None, sensitivity_relaxation=1.5, minimum_sensitivity=0):
		TrafficEngineer.__init__(self, aurora_network, all_interblock_paths)
		self.flow_pairs_to_ignore = ignore_flows
		self.update_frequency = update_frequency
		self.training_length = training_length
		self.reduce_multihop = reduce_multihop
		assert(sensitivity_relaxation >= 1)
		self.sensitivity_relaxation = sensitivity_relaxation
		self.minimum_sensitivity = minimum_sensitivity
		return

	def get_filename_param(self):
		return "dualhullte_r{}t{}c{}".format(self.update_frequency, self.training_length, self.num_k)

	def get_num_clusters_for_training(self):
		return self.num_k
	# returns a string description of this class
	def get_string(self):
		return "Dual Hull Method Sensitivity"

	def traffic_engineering_class(self):
		return TrafficEngineeringClass.ROBUST_TE


	# given an MLU, find the topology/routing weights such that the average hop count in the worst case
	# is minimized. This is an LP formulation that directly minimizes hop count, rather than previously done
	# which is implicit through the use of QP to minimize non-minimal path weights
	def _minimize_multihop_direct(self, logical_topology_adj_matrix, small_hull, large_hull, small_hull_mlu, large_hull_mlu):
		## using QP to reduce multihop reliance
		nblocks = self.aurora_network.get_num_blocks()
		num_tm = len(small_hull)

		model = Model("minimize multihop directly")
		model.setParam( 'OutputFlag', False )
		routing_weights_var = {}
		direct_hopcount = model.addVar(lb=0, ub=GRB.INFINITY, obj=1., vtype=GRB.CONTINUOUS, name="max_ahc")
		link_capacity_constraints_small_hull = [None] * num_tm
		link_capacity_constraints_large_hull = [None] * num_tm
		for tm_index in range(num_tm):
			link_capacity_constraints_small_hull[tm_index] = [None] * nblocks
			link_capacity_constraints_large_hull[tm_index] = [None] * nblocks
			for i in range(nblocks):
				link_capacity_constraints_small_hull[tm_index][i] = [None] * nblocks
				link_capacity_constraints_large_hull[tm_index][i] = [None] * nblocks
				for j in range(nblocks):
					if i != j:
						link_capacity_constraints_small_hull[tm_index][i][j] = LinExpr()
						link_capacity_constraints_large_hull[tm_index][i][j] = LinExpr()

		## setup the link constraints and all optization variables, 
		## which are the link counts and routing variables
		for i in range(nblocks):
			for j in range(nblocks):
				if (i != j):
					routing_weight_sum = LinExpr()
					for path in self.all_interblock_paths[i][j]:
						routing_weights_var[path] = model.addVar(lb=0, ub=GRB.INFINITY, obj=0., vtype=GRB.CONTINUOUS, name="w_{}".format(path))
						routing_weight_sum += routing_weights_var[path]
					model.addConstr(lhs=routing_weight_sum, sense=GRB.EQUAL, rhs=1)
		
		## add link capacity limit for all paths constraints
		for tm_index in range(num_tm):
			for i in range(nblocks):
				for j in range(nblocks):
					if (i != j):
						for path in self.all_interblock_paths[i][j]:
							path_len = len(path)
							curr_node = path[0]
							for next_node_index in range(1, path_len, 1):
								next_node = path[next_node_index]
								link_capacity_constraints_small_hull[tm_index][curr_node][next_node] += (routing_weights_var[path] * small_hull[tm_index][i][j])
								link_capacity_constraints_large_hull[tm_index][curr_node][next_node] += (routing_weights_var[path] * large_hull[tm_index][i][j])
								curr_node = next_node
		for tm_index in range(num_tm):
			for i in range(nblocks):
				for j in range(nblocks):
					if (i != j):
						capacity = min(self.aurora_network.get_link_capacity(i), self.aurora_network.get_link_capacity(j))
						model.addConstr(lhs=link_capacity_constraints_small_hull[tm_index][i][j], sense=GRB.LESS_EQUAL, rhs=small_hull_mlu * logical_topology_adj_matrix[i][j] * capacity)
						#model.addConstr(lhs=link_capacity_constraints_large_hull[tm_index][i][j], sense=GRB.LESS_EQUAL, rhs=large_hull_mlu * logical_topology_adj_matrix[i][j] * capacity)

		traffic_matrices_ahc = [None] * num_tm
		## maximum average hop count constraint
		for tm_index in range(num_tm):
			traffic_matrices_ahc[tm_index] = LinExpr()
			traffic_sum = float(sum([sum(x) for x in small_hull[tm_index]]))
			for i in range(nblocks):
				for j in range(nblocks):
					if i != j:
						traffic_matrices_ahc[tm_index] += (routing_weights_var[(i,j)] * (small_hull[tm_index][i][j]/traffic_sum))
			model.addConstr(lhs=direct_hopcount, sense=GRB.LESS_EQUAL, rhs=traffic_matrices_ahc[tm_index])

		# set up the objective function
		model.setObjective(direct_hopcount, GRB.MAXIMIZE)
		# start optimizing
		routing_weights = {}
		try: 
			model.optimize()
			print("ave hop count after routing is : {}".format(2. - direct_hopcount.x))
			for i in range(nblocks):
				for j in range(nblocks):
					if i != j:
						for path in self.all_interblock_paths[i][j]:
							routing_weights[path] = routing_weights_var[path].x
			#print adj_matrix
		except GurobiError as e:
			print ("MinimizeMultihop Direct Robust Multi Traffic TE: Error code " + str(e. errno ) + ": " + str(e))
		except AttributeError :
			print ("MinimizeMultihop Direct Robust Multi Traffic TE: Encountered an attribute error ")
		return routing_weights

	# given an MLU, find the topology/routing weights such that the average hop count in the worst case
	# is minimized. This is an LP formulation that directly minimizes hop count, rather than previously done
	# which is implicit through the use of QP to minimize non-minimal path weights
	def _minimize_multihop_qp(self, logical_topology_adj_matrix, small_hull, large_hull, small_hull_mlu, large_hull_mlu):
		## using QP to reduce multihop reliance
		nblocks = self.aurora_network.get_num_blocks()
		num_tm = len(small_hull)

		model = Model("minimize multihop directly")
		model.setParam( 'OutputFlag', False )
		routing_weights_var = {}
		qp_objective_function = QuadExpr()
		link_capacity_constraints_small_hull = [None] * num_tm
		link_capacity_constraints_large_hull = [None] * num_tm
		for tm_index in range(num_tm):
			link_capacity_constraints_small_hull[tm_index] = [None] * nblocks
			link_capacity_constraints_large_hull[tm_index] = [None] * nblocks
			for i in range(nblocks):
				link_capacity_constraints_small_hull[tm_index][i] = [None] * nblocks
				link_capacity_constraints_large_hull[tm_index][i] = [None] * nblocks
				for j in range(nblocks):
					if i != j:
						link_capacity_constraints_small_hull[tm_index][i][j] = LinExpr()
						link_capacity_constraints_large_hull[tm_index][i][j] = LinExpr()

		## setup the link constraints and all optization variables, 
		## which are the link counts and routing variables
		for i in range(nblocks):
			for j in range(nblocks):
				if (i != j):
					routing_weight_sum = LinExpr()
					for path in self.all_interblock_paths[i][j]:
						routing_weights_var[path] = model.addVar(lb=0, ub=GRB.INFINITY, obj=0., vtype=GRB.CONTINUOUS, name="w_{}".format(path))
						routing_weight_sum += routing_weights_var[path]
					model.addConstr(lhs=routing_weight_sum, sense=GRB.EQUAL, rhs=1)
		
		## add link capacity limit for all paths constraints
		for tm_index in range(num_tm):
			for i in range(nblocks):
				for j in range(nblocks):
					if (i != j):
						for path in self.all_interblock_paths[i][j]:
							path_len = len(path)
							curr_node = path[0]
							for next_node_index in range(1, path_len, 1):
								next_node = path[next_node_index]
								link_capacity_constraints_small_hull[tm_index][curr_node][next_node] += (routing_weights_var[path] * small_hull[tm_index][i][j])
								link_capacity_constraints_large_hull[tm_index][curr_node][next_node] += (routing_weights_var[path] * large_hull[tm_index][i][j])
								curr_node = next_node
		for tm_index in range(num_tm):
			for i in range(nblocks):
				for j in range(nblocks):
					if (i != j):
						capacity = min(self.aurora_network.get_link_capacity(i), self.aurora_network.get_link_capacity(j))
						model.addConstr(lhs=link_capacity_constraints_small_hull[tm_index][i][j], sense=GRB.LESS_EQUAL, rhs=small_hull_mlu * logical_topology_adj_matrix[i][j] * capacity)
						model.addConstr(lhs=link_capacity_constraints_large_hull[tm_index][i][j], sense=GRB.LESS_EQUAL, rhs=large_hull_mlu * logical_topology_adj_matrix[i][j] * capacity)

		## maximum average hop count constraint
		for i in range(nblocks):
			for j in range(nblocks):
				if i != j:
					for path in self.all_interblock_paths[i][j]:
						if self.reduce_multihop:
							if len(path) > 2:
								qp_objective_function += (routing_weights_var[path] * routing_weights_var[path])
						else:
							if len(path) == 2:
								qp_objective_function += (routing_weights_var[path] * routing_weights_var[path])

		# set up the objective function
		model.setObjective(qp_objective_function, GRB.MINIMIZE)
		# start optimizing
		routing_weights = {}
		try: 
			model.optimize()
			for i in range(nblocks):
				for j in range(nblocks):
					if i != j:
						for path in self.all_interblock_paths[i][j]:
							routing_weights[path] = routing_weights_var[path].x
			#print adj_matrix
		except GurobiError as e:
			print ("MinimizeMultihop Direct Robust Multi Traffic TE: Error code " + str(e. errno ) + ": " + str(e))
		except AttributeError :
			print ("MinimizeMultihop Direct Robust Multi Traffic TE: Encountered an attribute error ")
		return routing_weights


	## Currently does not involve another step that reduces hop count
	def _compute_path_weights_minimize_worstcase_mlu(self, logical_topology_adj_matrix, hull_traffic_matrices):
		nblocks = self.aurora_network.get_num_blocks()
		num_tm = len(hull_traffic_matrices)

		model = Model("Routing minimize MLU for all traffic matrices")
		model.setParam( 'OutputFlag', False )
		mlu = model.addVar(lb=0., ub=GRB.INFINITY, obj=0, vtype=GRB.CONTINUOUS, name="mlu")
		routing_weight_vars = {}
		link_capacity_constraints = [None] * num_tm
		for k in range(num_tm):
			link_capacity_constraints[k] = [None] * nblocks
			for i in range(nblocks):
				link_capacity_constraints[k][i] = [None] * nblocks
				for j in range(nblocks):
					if i != j:
						link_capacity_constraints[k][i][j] = LinExpr()
						#routing_weight_vars[path] = model.addVar(lb=0., ub=1., obj=1., vtype=GRB.CONTINUOUS, name="w_{},{}".format(i,j))
		for i in range(nblocks):
			for j in range(nblocks):
				if i != j:
					weight_constraint = LinExpr()
					for path in self.all_interblock_paths[i][j]:
						routing_weight_vars[path] = model.addVar(lb=0, ub=1., obj=1., vtype=GRB.CONTINUOUS, name="w_{}".format(path))
						weight_constraint += routing_weight_vars[path]
					model.addConstr(lhs=weight_constraint, sense=GRB.EQUAL, rhs=1.)


		# stage 1: setup all the optimization variables that are the routing weights
		# stage 1.5 : also adds the traffic flow satisfiability constraints
		for tm, tm_index in zip(hull_traffic_matrices, range(num_tm)):
			for i in range(nblocks):
				for j in range(nblocks):
					if i != j:
						for path in self.all_interblock_paths[i][j]:
							src = path[0]
							for path_hop in range(1, len(path), 1):
								dst = path[path_hop]
								link_capacity_constraints[tm_index][src][dst] += (routing_weight_vars[path] * tm[i][j])
								src = dst

		## stage 2 : add the link utilization constraints to the model
		for tm_index in range(num_tm):
			for i in range(nblocks):
				for j in range(nblocks):
					if i != j:
						lu_constraint = link_capacity_constraints[tm_index][i][j]
						link_capacity = min(self.aurora_network.get_link_capacity(i), self.aurora_network.get_link_capacity(j))
						model.addConstr(lhs=lu_constraint, sense=GRB.LESS_EQUAL, rhs=mlu * link_capacity * logical_topology_adj_matrix[i][j])

		# stage 3: set the objective
		model.setObjective(mlu, GRB.MINIMIZE)
		try: 
			model.optimize()
			routing_weights = {}
			for path in routing_weight_vars.keys():
				routing_weights[path] = routing_weight_vars[path].x
			return routing_weights, mlu.x
		except GurobiError as e:
			print ("Error code " + str(e. errno ) + ": " + str(e))
			return None
		except AttributeError :
			print ("Encountered an attribute error in Worst case MLU")
			return None

	def compute_path_weights(self, logical_topology_adj_matrix, large_hull, small_hull):
		## newly added for robustness
		nblocks = self.aurora_network.get_num_blocks()
		_, large_hull_mlu = self._compute_path_weights_minimize_worstcase_mlu(logical_topology_adj_matrix, large_hull)
		_, small_hull_mlu = self._compute_path_weights_minimize_worstcase_mlu(logical_topology_adj_matrix, small_hull)
		# compute the max TM
		## 
		routing_weights = self._minimize_multihop_qp(logical_topology_adj_matrix, small_hull, large_hull, small_hull_mlu, 2*large_hull_mlu)
		return routing_weights


